package gokv

import (
	"encoding/binary"
	"gitee.com/Cauchy_AQ/gokv/data"
	"github.com/elliotchance/orderedmap/v2"
	"sync"
	"sync/atomic"
)

// nonTransactionSeqNo 非事务标识常量
const nonTransactionSeqNo uint64 = 0

// txnFinKey 事务完成标识 key
var txnFinKey = []byte("txn-fin")

// WriteBatch 原子批量写，保证原子性
type WriteBatch struct {
	db            *DB
	mtx           *sync.Mutex
	option        *WriteBatchOptions
	pendingWrites *orderedmap.OrderedMap[string, *data.LogRecord] // 暂存用户写入数据，应为顺序结构
}

// NewWriteBatch 初始化 WriteBatch
func (d *DB) NewWriteBatch(opt *WriteBatchOptions) *WriteBatch {
	// TODO: 似乎 B+ 树需要控制，因为事务序列号
	return &WriteBatch{
		db:            d,
		mtx:           new(sync.Mutex),
		option:        opt,
		pendingWrites: orderedmap.NewOrderedMap[string, *data.LogRecord](),
	}
}

// Put 批量写数据
func (w *WriteBatch) Put(key []byte, value []byte) error {
	if len(key) == 0 {
		return ErrKeyIsEmpty
	}

	w.mtx.Lock()
	defer w.mtx.Unlock()

	logRecord := &data.LogRecord{
		Key:   key,
		Value: value,
	}
	w.pendingWrites.Set(string(key), logRecord)
	return nil
}

func (w *WriteBatch) Delete(key []byte) error {
	if len(key) == 0 {
		return ErrKeyIsEmpty
	}
	w.mtx.Lock()
	defer w.mtx.Unlock()

	logRecordPos := w.db.index.Get(key)
	// 数据不存在，而且现在执行删除，则直接清空前面的该数据记录
	if logRecordPos == nil {
		_, ok := w.pendingWrites.Get(string(key))
		if ok {
			_ = w.pendingWrites.Delete(string(key))
		}
		return nil
	}

	// 否则数据存在，则该动作存入待执行
	logRecord := &data.LogRecord{
		Key:  key,
		Type: data.LogRecordDeleted,
	}
	w.pendingWrites.Set(string(key), logRecord)

	return nil
}

// Commit 提交事务，批量数据持久化，并更新索引
func (w *WriteBatch) Commit() error {
	w.mtx.Lock()
	defer w.mtx.Unlock()

	if w.pendingWrites.Len() == 0 {
		return nil
	}

	if uint(w.pendingWrites.Len()) > w.option.MaxBatchNum {
		return ErrExceedMaxBatchNum
	}

	// 加锁保证事务提交串行化
	w.db.mtx.Lock()
	defer w.db.mtx.Unlock()

	// 获取当前最新事务序列号
	seqNo := atomic.AddUint64(&w.db.seqNo, 1)

	// 写入数据到数据文件
	positions := make(map[string]*data.LogRecordPos, w.pendingWrites.Len())
	for _, key := range w.pendingWrites.Keys() {
		record, _ := w.pendingWrites.Get(key)
		logRecordPos, err := w.db.appendLogRecord(&data.LogRecord{
			Key:   logRecordKeyWithSeq(record.Key, seqNo),
			Value: record.Value,
			Type:  record.Type,
		})
		if err != nil {
			return err
		}
		positions[string(record.Key)] = logRecordPos
	}

	// 写一条标识事务完成的数据
	finishedRecord := &data.LogRecord{
		Key:  logRecordKeyWithSeq(txnFinKey, seqNo),
		Type: data.LogRecordTxnFinished,
	}
	finRecordPos, err := w.db.appendLogRecord(finishedRecord)
	if err != nil {
		return err
	}
	if finRecordPos != nil {
		w.db.invalidDataSize += int64(finRecordPos.Size)
	}

	// 根据配置决定是否持久化
	if w.option.SyncWrites && w.db.activeFile != nil {
		if err := w.db.activeFile.Sync(); err != nil {
			return err
		}
	}

	// 更新内存索引
	for _, key := range w.pendingWrites.Keys() {
		record, _ := w.pendingWrites.Get(key)
		pos := positions[string(record.Key)]

		var oldPos *data.LogRecordPos
		if record.Type == data.LogRecordNormal {
			oldPos = w.db.index.Put(record.Key, pos)
		}

		if record.Type == data.LogRecordDeleted {
			oldPos, _ = w.db.index.Delete(record.Key)
		}

		if oldPos != nil {
			w.db.invalidDataSize += int64(oldPos.Size)
			if record.Type == data.LogRecordDeleted {
				w.db.invalidDataSize += int64(pos.Size)
			}
		}
	}

	// 清空缓存数据
	w.pendingWrites = orderedmap.NewOrderedMap[string, *data.LogRecord]()

	return nil
}

// logRecordKeyWithSeq SeqNo+key 编码
func logRecordKeyWithSeq(key []byte, seqNo uint64) []byte {
	seq := make([]byte, binary.MaxVarintLen64)
	n := binary.PutUvarint(seq[:], seqNo)

	encKey := make([]byte, n+len(key))
	copy(encKey[:n], seq[:n])
	copy(encKey[n:], key)
	return encKey
}

// parseLogRecordKey 解决 LogRecord 的 key，获取实际 key 和事务序列号
func parseLogRecordKey(key []byte) ([]byte, uint64) {
	seqNo, n := binary.Uvarint(key)
	realKey := key[n:]
	return realKey, seqNo
}
